﻿#include "DILUPrecond.hpp"
#include "csrMatrix.hpp"

template <typename Cmpt>
UNAP::DILUPrecond<Cmpt>::DILUPrecond(const CSRMatrix<Cmpt> &A)
    : Preconditioner<Cmpt>(A.getCommunicator()), rD_(A.diag()) {
  const label rowSize = A.lRows();
  const label low = A.beginRow();
  const label high = low + rowSize;
  const label *ptr = A.ptr();
  const label *offDiagStart = A.offDiagStart().data();
  const label *ind = A.ind();
  const Cmpt *val = A.val();
  Array<Triplet<Cmpt> > upper, lower;
  upper.reserve((A.lnnz() - A.offDiagNnz()) / 2);
  lower.reserve((A.lnnz() - A.offDiagNnz()) / 2);
  for (label row = 0; row < rowSize; row++) {
    for (label i = ptr[row]; i < offDiagStart[row]; i++) {
      const label col = ind[i];
      const Cmpt value = val[i];
      if (col >= low && col < high) {
        if (row + low < col)
          upper.push_back(Triplet<Cmpt>(row, col, value));
        else if (row + low > col)
          lower.push_back(Triplet<Cmpt>(col, row, value));
      }
    }
  }

  nFaces_ = upper.size();
  uPtr_ = new label[nFaces_];
  lPtr_ = new label[nFaces_];
  losortPtr_ = new label[nFaces_];
  upperPtr_ = new Cmpt[nFaces_];
  lowerPtr_ = new Cmpt[nFaces_];

  for (label i = 0; i < nFaces_; i++) {
    uPtr_[i] = upper[i].c_ - low;
    lPtr_[i] = upper[i].r_;
    upperPtr_[i] = upper[i].v_;
  }

  {
    Array<label> nNbrOfFace(rowSize);
    forAll(nbrI, nFaces_) { nNbrOfFace[uPtr_[nbrI]]++; }
    label **cellNbrFaces = new label *[rowSize];
    forAll(cellI, rowSize) {
      cellNbrFaces[cellI] = new label[nNbrOfFace[cellI]];
    }
    std::fill(nNbrOfFace.begin(), nNbrOfFace.end(), 0);
    forAll(nbrI, nFaces_) {
      cellNbrFaces[uPtr_[nbrI]][nNbrOfFace[uPtr_[nbrI]]] = nbrI;
      nNbrOfFace[uPtr_[nbrI]]++;
    }
    label lstI = 0;
    forAll(cellI, rowSize) {
      forAll(curNbrI, nNbrOfFace[cellI]) {
        losortPtr_[lstI] = cellNbrFaces[cellI][curNbrI];
        lowerPtr_[losortPtr_[lstI]] = lower[lstI].v_;
        lstI++;
      }
    }

    forAll(cellI, rowSize) delete[] cellNbrFaces[cellI];
    delete[] cellNbrFaces;
  }

  ownData_ = true;
  calcReciprocalD(rD_);
}

template class UNAP::DILUPrecond<scalar>;
